from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization


def load_public_key_from_file(file_path):
    """从PEM文件加载公钥"""
    with open(file_path, "rb") as key_file:
        public_key = serialization.load_pem_public_key(
            key_file.read(), backend=default_backend()
        )
    return public_key


def rsa_encrypt(public_key, plaintext):
    """使用公钥加密数据"""
    # 确保明文是字节类型
    if isinstance(plaintext, str):
        plaintext = plaintext.encode("utf-8")

    ciphertext = public_key.encrypt(
        plaintext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return ciphertext


if __name__ == "__main__":
    # 加载本地公钥文件
    public_key_path = "public_key.pem"  # 公钥文件路径
    try:
        public_key = load_public_key_from_file(public_key_path)
        print(f"已成功从 {public_key_path} 加载公钥")
    except Exception as e:
        print(f"加载公钥失败: {e}")
        exit(1)

    # 要加密的原始数据
    original_data = "这是需要加密的数据，可以包含中文和特殊字符！@#$%^&*()"
    print(f"\n原始数据: {original_data}")

    # 使用公钥加密
    encrypted_data = rsa_encrypt(public_key, original_data)

    # 打印加密后的数据（以十六进制格式，便于查看和传输）
    print(f"\n加密后的数据 (十六进制): {encrypted_data.hex()}")

    # 也可以打印原始字节数据（可读性较差）
    print(f"加密后的数据 (字节): {encrypted_data}")
